# 模糊去噪
import random

import cv2
import numpy as np
import os

'''
1、生成一个包含噪声的灰度图并保存到data文件夹下
'''
img = cv2.imread(os.path.join('../data', 'xm.png'))

# 确保data文件夹存在，如果不存在就创建
save_dir = '../data'
if not os.path.exists(save_dir):
    os.makedirs(save_dir)
    print(f'创建文件夹：{save_dir}')
# 改变灰度
# cvtColor用于颜色空间转换 COLOR_BGR2GRAY表示从BGR格式里面转换为灰度格式
img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)


# 为灰度图片加入噪点
def add_noise_salt(image, prob=0.05):
    """
    :param image: 输入的图像（灰度图）
    :param prob: 噪声概率值，值越大噪点越多
    :return: 返回添加的噪声图片
    """
    output = np.copy(image)
    # 计算添加多少个白色的噪点（椒盐）
    thres = 1 - prob
    for i in range(image.shape[0]):
        for j in range(image.shape[1]):
            rdm = np.random.random()  # 生成0 -1 之间的随机数
            if rdm < prob:
                output[i][j] = 0  # 黑色 椒噪点
            elif rdm > thres:
                output[i][j] = 255  # 盐噪声（白色）
    return output


# 生成灰度图并包含噪点图像
salt_img = add_noise_salt(img_gray)
# 保存椒盐噪点的图片
sp_save_path= os.path.join(save_dir,"xm_gray.png")
cv2.imwrite(sp_save_path,salt_img)
print(f'椒盐噪声图片已保存到：{sp_save_path}')

cv2.imshow('img', img)
cv2.imshow('img_gray', img_gray)
cv2.imshow('salt_img', salt_img)

cv2.waitKey(0)
